//
//  KeyStoreManager.swift
//  KeyStorePlugin
//
//  Created by bzb on 2023/9/25.
//  Copyright © 2023 CocoaPods. All rights reserved.
//

import Foundation
import Security
import CommonCrypto

// Singleton instance
public let _singletonInstance = KeyStoreManager()


// Constants
public var kRSASwiftGeneratorKeyType = kSecAttrKeyTypeRSA
public var kRSASwiftGeneratorKeySize = 1024
public var kRSASwiftGeneratorCypheredBufferSize = 1024
public var kRSASwiftGeneratorSecPadding: SecPadding = .PKCS1
public var MAX_ENCRYPT_BLOCK = 117
public var MAX_DECRYPT_BLOCK = 128


public enum CryptoException: Error {
    case unknownError
    case duplicateFoundWhileTryingToCreateKey
    case keyNotFound
    case authFailed
    case unableToAddPublicKeyToKeyChain
    case wrongInputDataFormat
    case unableToEncrypt
    case unableToDecrypt
    case unableToSignData
    case unableToVerifySignedData
    case unableToPerformHashOfData
    case unableToGenerateAccessControlWithGivenSecurity
    case outOfMemory
}

public class KeyStoreManager: NSObject {
    
   /** Shared instance */
   public class var instance: KeyStoreManager {
        return _singletonInstance
    }
    
    /// RSA密钥生成, 返回Base64编码后的公钥
    func generateRSA(_ alias: String) -> Data? {
        // 判断是否已存在
        if keyPairExists(alias) {
            let data = self.getPublicKeyData(alias)
            print("generateRSA: alias=\(alias) has exists, pubKey: \(data?.base64EncodedString())")
            return data
        } else {
            print("generateRSA: alias=\(alias) not exists, starting create...")
        }
        
        let ref = PKCSLoadRSAKey(kSecAttrKeyClassPublic, alias)
        if ref == nil {
            print("generateRSA: 本地不存在此密钥")
//            let isSuccess = PCKSGenerateKeyPair(alias, alias, 256)
//            print(isSuccess ? "密钥创建成功" : "密钥创建失败")
            let pubKey = generateRSAPair(alias)
            print(pubKey != nil ? "generateRSA: 密钥创建成功" : "generateRSA: 密钥创建失败")
            return pubKey
        } else {
            let pubKey = getPublicKeyData(alias)
            print("generateRSA: 本地已存在此密钥")
            return pubKey
        }
    }
    
    /// RSA密钥生成, , 返回公钥的 N 值
    func generateRSAByModulus(_ alias: String) -> String? {
        // 判断是否已存在
        if keyPairExists(alias) {
            let nData = self.getPubKeyNEValueFromPublicKey(alias)
//            print("generateRSA: alias=\(alias) has exists, pubKey N 值: \(nData)")
            return nData
        } else {
            print("generateRSA: alias=\(alias) not exists, starting create...")
        }
        
        let ref = PKCSLoadRSAKey(kSecAttrKeyClassPublic, alias)
        if ref == nil {
            print("generateRSA: 本地不存在此密钥")
//            let isSuccess = PCKSGenerateKeyPair(alias, alias, 256)
//            print(isSuccess ? "密钥创建成功" : "密钥创建失败")
            let pubKey = generateRSAPair(alias)
            let nData = getPubKeyNEValueFromPublicKey(alias)
            print(pubKey != nil ? "generateRSA: 密钥创建成功" : "generateRSA: 密钥创建失败")
            return nData
        } else {
            let nData = getPubKeyNEValueFromPublicKey(alias)
            print("generateRSA: 本地已存在此密钥")
            return nData
        }
    }
    
    /// RSA加密
    func encryptRSA(_ alias: String, _ content: String) -> String? {
        let isExists = keyPairExists(alias)
        print("encryptRSA: 密钥 \(alias) isExists=\(isExists)")
        
        if !isExists {
            return nil
        }
        
        let pubKeyRef = PKCSLoadRSAKey(kSecAttrKeyClassPublic, alias)
        if pubKeyRef != nil {
            // 加密
            let result = PKCSEncryptRSA(content, pubKeyRef?.takeRetainedValue())
            print(result ?? "encryptRSA: 加密结果为空")
            return result
        } else {
            print("encryptRSA: 公钥未获取到")
            return nil
        }
    }
    
    /// RSA解密
    func decryptRSA(_ alias: String, _ decryptContent: String) -> String? {
        let isExists = keyPairExists(alias)
        print("decryptRSA: 密钥 \(alias) isExists=\(isExists)")
        
        if !isExists {
            return nil
        }
        
        let priKeyRef = PKCSLoadRSAKey(kSecAttrKeyClassPrivate, alias)
        
        // 解密
        if priKeyRef != nil && decryptContent != "" {
            let result = PKCSDecryptRSA(decryptContent, priKeyRef?.takeRetainedValue())
            print(result ?? "decryptRSA: 解密结果为空")
            return result
        } else {
            print("decryptRSA: 私钥未获取到, 或需解密内容为空")
            return nil
        }
    }
    
    /**
     * 加密后端分段加密的白盒密钥大数据(分段加密)
     */
    func encryptBigRSA(_ alias: String, _ decryptContent: String) -> String? {
        let isExists = keyPairExists(alias)
        if !isExists {
            return nil
        }
        
        let pubKeyRef = PKCSLoadRSAKey(kSecAttrKeyClassPublic, alias)
        let pubKey = pubKeyRef?.takeRetainedValue()
        
        let inputLength = decryptContent.count
        var i = 0
        var offset = 0
        var out: String = ""
        
        while (inputLength - offset > 0) {
            var temp: String
            let start = decryptContent.index(decryptContent.startIndex, offsetBy: offset)
            if (inputLength - offset > MAX_ENCRYPT_BLOCK) {
                let end = decryptContent.index(decryptContent.startIndex, offsetBy: offset + MAX_ENCRYPT_BLOCK)
                temp = String(decryptContent[start ..< end])
            } else {
                let end = decryptContent.index(decryptContent.startIndex, offsetBy: inputLength)
                temp = String(decryptContent[start ..< end])
            }
            
            let cache = PKCSEncryptRSA(temp, pubKey) ?? ""
            out += cache
            i += 1
            offset = i * MAX_ENCRYPT_BLOCK
        }

        return out
    }
    
    /**
     * 解密后端分段加密的白盒密钥大数据(分段解密)
     */
    func decryptBigRSA(_ alias: String, _ decryptContent: String) -> String? {
        let isExists = keyPairExists(alias)
        if !isExists {
            return nil
        }
        
        let priKeyRef = PKCSLoadRSAKey(kSecAttrKeyClassPrivate, alias)
        let priKey = priKeyRef?.takeRetainedValue()
        
        let blockSize = SecKeyGetBlockSize(priKey!)
        let encryptedData = Data(base64Encoded: decryptContent)
        var inputLength = encryptedData!.count
        var offset = 0
        var decryptedData = Data()

        while inputLength - offset > 0 {
            var cacheLength = 0
            if inputLength - offset > blockSize {
                cacheLength = blockSize
            } else {
                cacheLength = inputLength - offset
            }

            var cache = [UInt8](repeating: 0, count: blockSize)
            encryptedData!.copyBytes(to: &cache, from: offset..<offset+cacheLength)

            var decryptedLength = blockSize
            var decryptedBytes = [UInt8](repeating: 0, count: decryptedLength)

            let status = SecKeyDecrypt(priKey!, .PKCS1, &cache, cacheLength, &decryptedBytes, &decryptedLength)
            if status == errSecSuccess {
                decryptedData.append(decryptedBytes, count: decryptedLength)
            } else {
                print(status)
//                throw NSError(domain: NSOSStatusErrorDomain, code: Int(status), userInfo: nil)
            }

            offset += cacheLength
        }

        return String(data: decryptedData, encoding: .utf8)
    }
    
    func decryptBig(_ alias: String, _ decryptContent: String) -> String? {
        let isExists = keyPairExists(alias)
        if !isExists {
            return nil
        }
        
        let priKeyRef = PKCSLoadRSAKey(kSecAttrKeyClassPrivate, alias)
        let priKey = priKeyRef?.takeRetainedValue()
        
        let inputLength = decryptContent.count
        var i = 0
        var offset = 0
        var out: String = ""
        
        while (inputLength - offset > 0) {
            var temp: String
            let start = decryptContent.index(decryptContent.startIndex, offsetBy: offset)
            if (inputLength - offset > MAX_DECRYPT_BLOCK) {
                let end = decryptContent.index(decryptContent.startIndex, offsetBy: offset + MAX_DECRYPT_BLOCK)
                temp = String(decryptContent[start ..< end])
            } else {
                let end = decryptContent.index(decryptContent.startIndex, offsetBy: inputLength)
                temp = String(decryptContent[start ..< end])
            }
            
            let cache = PKCSDecryptRSA(temp, priKey) ?? ""
            out += cache
            i += 1
            offset = i * MAX_DECRYPT_BLOCK
        }

        return out
    }

    /// 创建一对RSA密钥, 返回公钥的Base64编码字符串
    /// 本方法不做是否存在检验, 如本地已有, 会覆盖创建
    private func generateRSAPair(_ alias: String) -> Data? {
         // private key parameters
         let privateKeyParams: [String: AnyObject] = [
             kSecAttrIsPermanent as String: true as AnyObject,
             kSecAttrApplicationTag as String: alias as AnyObject
         ]

         // public key parameters
         let publicKeyParams: [String: AnyObject] = [
             kSecAttrIsPermanent as String: true as AnyObject,
             kSecAttrApplicationTag as String: alias as AnyObject
         ]

         // global parameters for our key generation
         let parameters: [String: AnyObject] = [
             kSecAttrKeyType as String:          kRSASwiftGeneratorKeyType,
             kSecAttrKeySizeInBits as String:    kRSASwiftGeneratorKeySize as AnyObject,
             kSecPublicKeyAttrs as String:       publicKeyParams as AnyObject,
             kSecPrivateKeyAttrs as String:      privateKeyParams as AnyObject,
             ]

        // 创建密钥对
        var pubKey, privKey: SecKey?
        var data: Data?

        let status = SecKeyGeneratePair(parameters as CFDictionary, &pubKey, &privKey)
        if status == errSecSuccess {
          data = getPublicKeyData(alias)
//          print("createSecureKeyPair success, pubKeyBase64: \(data?.base64EncodedString())")
//          print("")
//          print("createSecureKeyPair success, pubKey: \(pubKey)")
        } else {
          var error = CryptoException.unknownError
          switch (status) {
            case errSecDuplicateItem: error = .duplicateFoundWhileTryingToCreateKey
            case errSecItemNotFound: error = .keyNotFound
            case errSecAuthFailed: error = .authFailed
            default: break
          }
          print("createSecureKeyPair error: \(error)")
        }
        return data
    }

    /// 获取公钥的 N & E 值的 16 进制字符串格式, 字符总长度 266, 前 256 是 N 值, 后10是E值
    func getPubKeyNEValueFromPublicKey(_ alias: String) -> String? {
        let publicKey = KeyStoreManager.instance.getPublicKeyReference(alias)
        if publicKey == nil {
            return nil
        }
        
        var error: Unmanaged<CFError>?
        guard let publicKeyData = SecKeyCopyExternalRepresentation(publicKey!, &error) as Data? else {
            print("Error while retrieving public key data: \(error.debugDescription)")
            return nil
        }

        let modulusTag: [UInt8] = [
            0x02, 0x81, 0x81, 0x00
        ]

        guard let index = publicKeyData.range(of: Data(modulusTag), options: .backwards, in: 0..<publicKeyData.count) else {
            print("Modulus tag not found")
            return nil
        }

        let modulusData = publicKeyData.subdata(in: index.upperBound..<publicKeyData.endIndex)
        return modulusData.map { String(format: "%02X", $0) }.joined()
    }
    
    /// 获取公钥数据
    public func getPublicKeyData(_ alias: String) -> Data? {
        let parameters = [
            kSecClass as String: kSecClassKey,
            kSecAttrKeyType as String: kSecAttrKeyTypeRSA,
            kSecAttrApplicationTag as String: alias,
            kSecAttrKeyClass as String: kSecAttrKeyClassPublic,
            kSecReturnData as String: true
            ] as [String : Any]
        var data: AnyObject?
        let status = SecItemCopyMatching(parameters as CFDictionary, &data)
        if status == errSecSuccess {
            return data as? Data
        } else { return nil }
    }

    /// 获取公钥对象引用
   public func getPublicKeyReference(_ alias: String) -> SecKey? {
        let parameters = [
            kSecClass as String: kSecClassKey,
            kSecAttrKeyType as String: kSecAttrKeyTypeRSA,
            kSecAttrApplicationTag as String: alias,
            kSecAttrKeyClass as String: kSecAttrKeyClassPublic,
            kSecReturnRef as String: true,
            ] as [String : Any]
        var ref: AnyObject?
        let status = SecItemCopyMatching(parameters as CFDictionary, &ref)
        if status == errSecSuccess { return ref as! SecKey? } else { return nil }
    }


    /// 获取私钥对象引用
   public func getPrivateKeyReference(_ alias: String) -> SecKey? {
        let parameters = [
            kSecClass as String: kSecClassKey,
            kSecAttrKeyClass as String: kSecAttrKeyClassPrivate,
            kSecAttrApplicationTag as String: alias,
            kSecReturnRef as String: true,
            ] as [String : Any]
        var ref: AnyObject?
        let status = SecItemCopyMatching(parameters as CFDictionary, &ref)
        if status == errSecSuccess { return ref as! SecKey? } else { return nil }
    }

    /// 判断密钥是否存在于 KeyChain
    public func keyPairExists(_ alias: String) -> Bool {
        return self.getPublicKeyData(alias) != nil
    }

    /// 删除密钥
   public func deleteSecureKeyPair(_ alias: String) -> Bool {
        // private query dictionary
        let deleteQuery = [
            kSecClass as String: kSecClassKey,
            kSecAttrApplicationTag as String: alias,
            ] as [String : Any]

       let status = SecItemDelete(deleteQuery as CFDictionary)
       return status == errSecSuccess
    }


    // MARK: - Cypher and decypher methods
    public func encryptMessageWithPublicKey(_ alias: String, _ message: String, completion: @escaping (_ success: Bool, _ data: Data?, _ error: CryptoException?) -> Void) {
         DispatchQueue.global(qos: DispatchQoS.QoSClass.default).async { () -> Void in
             
             if let publicKeyRef = self.getPublicKeyReference(alias) {
                 guard let messageData = message.data(using: String.Encoding.utf8) else {
                     completion(false, nil, .wrongInputDataFormat)
                     return
                 }
                 
                 let plainText = (messageData as NSData).bytes.bindMemory(to: UInt8.self, capacity: messageData.count)
                 let plainTextLen = messageData.count
                 let cipherData = Data(count: SecKeyGetBlockSize(publicKeyRef))
                 let cipherText: UnsafeMutablePointer<UInt8> = cipherData.withUnsafeBytes { $0.load(as: UnsafeMutablePointer<UInt8>.self) }
                 var cipherTextLen = cipherData.count
                 let status = SecKeyEncrypt(publicKeyRef, .PKCS1, plainText, plainTextLen, cipherText, &cipherTextLen)
                 
                 // analyze results and call the completion in main thread
                 DispatchQueue.main.async(execute: { () -> Void in
                     completion(status == errSecSuccess, cipherData, status == errSecSuccess ? nil : .unableToEncrypt)
                     cipherText.deinitialize(count: cipherTextLen)
                 })
                 return
             } else { DispatchQueue.main.async(execute: { completion(false, nil, .keyNotFound) }) }
         }
     }

    /// 使用公钥加密数据
   public func encryptMessageWithPublicKey(_ alias: String, _ content: String) -> String? {
       if let publicKeyRef = self.getPublicKeyReference(alias) {
           guard let encryptedData = content.data(using: String.Encoding.utf8) else {
               print("encryptMessageWithPublicKey error: .wrongInputDataFormat")
               return nil
           }
           let plainText = (encryptedData as NSData).bytes.bindMemory(to: UInt8.self, capacity: encryptedData.count)
           let plainTextLen = encryptedData.count
           let cipherData = Data(count: SecKeyGetBlockSize(publicKeyRef))
           let cipherText: UnsafeMutablePointer<UInt8> = cipherData.withUnsafeBytes { $0.load(as: UnsafeMutablePointer<UInt8>.self) }
           var cipherTextLen = cipherData.count
           
           DispatchQueue.global(qos: DispatchQoS.QoSClass.default).async { () -> Void in
               let status = SecKeyEncrypt(publicKeyRef, .PKCS1, plainText, plainTextLen, cipherText, &cipherTextLen)
               DispatchQueue.main.async(execute: { () -> Void in
                   // 根据状态判断是否加密成功
                   if status != errSecSuccess {
                       print("encryptMessageWithPublicKey fail: .unableToEncrypt")
                   } else {
                       print("encryptMessageWithPublicKey success: \(cipherData.base64EncodedString())")
                   }
                   cipherText.deinitialize(count: cipherTextLen)
                   
               })
           }
           
           return nil
           
//           var error: Unmanaged<CFError>?
//           if content == "" {
//               print("encryptMessageWithPublicKey error: .wrongInputDataFormat")
//               return nil
//           }
//           let encryptedData = content.data(using: .utf8)! as CFData
//
//           let data = SecKeyCreateEncryptedData(publicKeyRef, SecKeyAlgorithm.rsaEncryptionPKCS1, encryptedData, &error)
//           return String(data: data!)
       } else {
           print("encryptMessageWithPublicKey error: .keyNotFound")
           return nil
       }
    }

    /// 使用私钥解密数据
   public func decryptMessageWithPrivateKey(_ alias: String, _ content: String) -> String? {
       if let privateKeyRef = self.getPrivateKeyReference(alias) {
           guard let decryptedData = content.data(using: String.Encoding.utf8) else {
               print("encryptMessageWithPublicKey error: .wrongInputDataFormat")
               return nil
           }
           
           let encryptedText = (decryptedData as NSData).bytes.bindMemory(to: UInt8.self, capacity: decryptedData.count)
           let encryptedTextLen = decryptedData.count
           var plainData = Data(count: kRSASwiftGeneratorCypheredBufferSize)
           let plainText: UnsafeMutablePointer<UInt8> = plainData.withUnsafeBytes { $0.load(as: UnsafeMutablePointer<UInt8>.self) }
           var plainTextLen = plainData.count

           let status = SecKeyDecrypt(privateKeyRef, .PKCS1, encryptedText, encryptedTextLen, plainText, &plainTextLen)

           // analyze results and call the completion in main thread
           if status == errSecSuccess {
               plainData.count = plainTextLen
               // Generate and return result string
               if let string = NSString(data: plainData as Data, encoding: String.Encoding.utf8.rawValue) as String? {
                   return string
               } else {
                   print("decryptMessageWithPrivateKey fail: .unableToDecrypt")
               }
           } else {
               print("decryptMessageWithPrivateKey status fail: .unableToDecrypt")
           }
           plainText.deinitialize(count: plainTextLen)
           return nil
       } else {
           print("decryptMessageWithPrivateKey error: .keyNotFound")
           return nil
       }
    }
}



